#include "ws_connection.h"

WSConnection::WSConnection(boost::asio::io_context& ioc, boost::asio::ssl::context& ssl_ctx)
    : ioc_(ioc), ssl_ctx_(ssl_ctx), resolver_(ioc) {}

WSConnection::~WSConnection() {
    try {
        close();
    } catch (...) {
        // ignore any exceptions during destruction
    }
}

void WSConnection::connect(const std::string& host, const std::string& port, const std::string& path) {
    host_ = host;
    path_ = path;
    resolver_.async_resolve(
        host, port,
        std::bind(&WSConnection::onResolve, shared_from_this(),
                  std::placeholders::_1, std::placeholders::_2));
}

void WSConnection::onResolve(const boost::system::error_code& ec,
                             boost::asio::ip::tcp::resolver::results_type results) {
    if (ec) {
        if (onError_) onError_("Resolve failed: " + ec.message());
        return;
    }
    ws_ = std::make_unique<boost::beast::websocket::stream<
        boost::beast::ssl_stream<boost::asio::ip::tcp::socket>>>(ioc_, ssl_ctx_);
    boost::asio::async_connect(
        ws_->next_layer().next_layer(),
        results.begin(),
        results.end(),
        std::bind(&WSConnection::onConnect, shared_from_this(), std::placeholders::_1));
}

void WSConnection::onConnect(const boost::system::error_code& ec) {
    if (ec) {
        if (onError_) onError_("TCP connect failed: " + ec.message());
        return;
    }
    ws_->next_layer().async_handshake(
        boost::asio::ssl::stream_base::client,
        std::bind(&WSConnection::onSSLHandshake, shared_from_this(), std::placeholders::_1));
}

void WSConnection::onSSLHandshake(const boost::system::error_code& ec) {
    if (ec) {
        if (onError_) onError_("SSL handshake failed: " + ec.message());
        return;
    }
    ws_->async_handshake(
        host_, path_,
        std::bind(&WSConnection::onWSHandshake, shared_from_this(), std::placeholders::_1));
}

void WSConnection::onWSHandshake(const boost::system::error_code& ec) {
    if (ec) {
        if (onError_) onError_("WebSocket handshake failed: " + ec.message());
        return;
    }
    if (onConnect_) onConnect_();
    doRead();
}

void WSConnection::doRead() {
    ws_->async_read(
        buffer_,
        std::bind(&WSConnection::onRead, shared_from_this(),
                  std::placeholders::_1, std::placeholders::_2));
}

void WSConnection::onRead(const boost::system::error_code& ec, std::size_t bytes_transferred) {
    if (ec) {
        if (ec != boost::beast::websocket::error::closed && onError_) {
            onError_("Read failed: " + ec.message());
        }
        return;
    }

    std::string msg = boost::beast::buffers_to_string(buffer_.data());
    buffer_.consume(bytes_transferred);

    if (onMessage_) onMessage_(msg);
    doRead();
}

void WSConnection::send(const std::string& message) {
    if (!ws_ || !ws_->is_open()) return;

    boost::beast::error_code ec;
    ws_->write(boost::asio::buffer(message), ec);
    if (ec && onError_) onError_("Write failed: " + ec.message());
}

void WSConnection::close() {
    if (ws_ && ws_->is_open()) {
        ws_->async_close(
            boost::beast::websocket::close_code::normal,
            std::bind(&WSConnection::onClose, shared_from_this(), std::placeholders::_1));
            
    }
}

void WSConnection::onClose(const boost::system::error_code& ec) {
    if (ec && onError_) onError_("Close failed: " + ec.message());
}

bool WSConnection::isOpen() const {
    return ws_ && ws_->is_open();
}
